/*
 *  Copyright (c) 2006 Alessandro Pilotti
 *
 *  Permission is hereby granted, free of charge, to any person
 *  obtaining a copy of this software and associated documentation
 *  files (the "Software"), to deal in the Software without
 *  restriction, including without limitation the rights to use,
 *  copy, modify, merge, publish, distribute, sublicense, and/or sell
 *  copies of the Software, and to permit persons to whom the
 *  Software is furnished to do so, subject to the following
 *  conditions:
 *  
 *  The above copyright notice and this permission notice shall be
 *  included in all copies or substantial portions of the Software.
 *  
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 *  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 *  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 *  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 *  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 *  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 *  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 *  OTHER DEALINGS IN THE SOFTWARE.
 */

using System;
using System.Collections.Generic;
using System.Text;
using System.Runtime.InteropServices;
using System.ComponentModel;

namespace AlexPilotti.DPAPI
{
    public class DPAPIWrapper
    {
        #region PInvoke private methods

        [DllImport("Crypt32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        private static extern bool CryptProtectData(
          ref DATA_BLOB pDataIn,
          String szDataDescr,
          ref DATA_BLOB pOptionalEntropy,
          IntPtr pvReserved,
          ref CRYPTPROTECT_PROMPTSTRUCT pPromptStruct,
          int dwFlags,
          ref DATA_BLOB pDataOut);

        [DllImport("crypt32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        private static extern bool CryptUnprotectData(
               ref DATA_BLOB pCipherText,
               ref string pszDescription,
               ref DATA_BLOB pEntropy,
               IntPtr pReserved,
               ref CRYPTPROTECT_PROMPTSTRUCT pPrompt,
               int dwFlags,
               ref DATA_BLOB pPlainText); 

/*
        [DllImport("kernel32.dll", CharSet = CharSet.Auto)]
        private static extern int FormatMessage(
          int dwFlags,
          IntPtr lpSource,
          int dwMessageId,
          int dwLanguageId,
          ref String lpBuffer,
          int nSize,
          IntPtr Arguments);
 */ 

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        internal struct DATA_BLOB
        {
            public int cbData;
            public IntPtr pbData;
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        internal struct CRYPTPROTECT_PROMPTSTRUCT
        {
            public int cbSize;
            public int dwPromptFlags;
            public IntPtr hwndApp;
            public String szPrompt;
        }

        #endregion

        private const int CRYPTPROTECT_UI_FORBIDDEN = 0x1;
        private const int CRYPTPROTECT_LOCAL_MACHINE = 0x4;

        public enum Store
        {
            /// <summary>
            /// <c>USE_MACHINE_STORE</c>
            /// </summary>
            Machine = 1,
            /// <summary>
            /// <c>USE_USER_STORE</c>
            /// </summary>
            User            
        };

        private static void InitPromptstruct(ref CRYPTPROTECT_PROMPTSTRUCT ps)
        {
            ps.cbSize = Marshal.SizeOf(typeof(CRYPTPROTECT_PROMPTSTRUCT));
            ps.dwPromptFlags = 0;
            ps.hwndApp = IntPtr.Zero;
            ps.szPrompt = null;
        }

/*
        private static String GetErrorMessage(int errorCode)
        {
            int FORMAT_MESSAGE_ALLOCATE_BUFFER = 0x00000100;
            int FORMAT_MESSAGE_IGNORE_INSERTS = 0x00000200;
            int FORMAT_MESSAGE_FROM_SYSTEM = 0x00001000;
            int messageSize = 255;
            String lpMsgBuf = "";
            int dwFlags = FORMAT_MESSAGE_ALLOCATE_BUFFER |
                          FORMAT_MESSAGE_FROM_SYSTEM |
                          FORMAT_MESSAGE_IGNORE_INSERTS;

            int retVal = FormatMessage(dwFlags, IntPtr.Zero, errorCode, 0, ref lpMsgBuf, messageSize, IntPtr.Zero);
            if(0 == retVal)
            {
                throw new DPAPIException("Failed to format message for error code " + errorCode + ". ");
            }
            return lpMsgBuf;
        }
*/ 

        public static byte[] Encrypt(byte[] plainText)
        {
            return Encrypt(plainText, null);
        }

        public static byte[] Encrypt(byte[] plainText, byte[] optionalEntropy)
        {
            return Encrypt(plainText, optionalEntropy, Store.User);
        }

        /// <summary>
        /// Encrypt data using <c>CryptProtectData</c>. Returns the encrypted data.
        /// </summary>
        /// <param name="plainText"></param>
        /// <param name="optionalEntropy"></param>
        /// <param name="store"></param>
        /// <returns></returns>
        public static byte[] Encrypt(byte[] plainText, byte[] optionalEntropy, Store store)
        {
            bool retVal = false;
            DATA_BLOB plainTextBlob = new DATA_BLOB();
            DATA_BLOB cipherTextBlob = new DATA_BLOB();
            DATA_BLOB entropyBlob = new DATA_BLOB();
            CRYPTPROTECT_PROMPTSTRUCT prompt = new CRYPTPROTECT_PROMPTSTRUCT();
            InitPromptstruct(ref prompt);
            int dwFlags;
            try
            {
                int bytesSize = plainText.Length;
                plainTextBlob.pbData = Marshal.AllocHGlobal(bytesSize);
                if (plainTextBlob.pbData == IntPtr.Zero)
                    throw new DPAPIException("Unable to allocate plaintext buffer");

                plainTextBlob.cbData = bytesSize;
                Marshal.Copy(plainText, 0, plainTextBlob.pbData, bytesSize);

                if (store == Store.Machine)
                {
                    dwFlags = CRYPTPROTECT_LOCAL_MACHINE | CRYPTPROTECT_UI_FORBIDDEN;

                    if (optionalEntropy == null)
                        optionalEntropy = new byte[0];

                    bytesSize = optionalEntropy.Length;
                    entropyBlob.pbData = Marshal.AllocHGlobal(optionalEntropy.Length);
                    if (entropyBlob.pbData == IntPtr.Zero)
                        throw new DPAPIException("Unable to allocate entropy data buffer");

                    Marshal.Copy(optionalEntropy, 0, entropyBlob.pbData, bytesSize);
                    entropyBlob.cbData = bytesSize;
                }
                else
                    dwFlags = CRYPTPROTECT_UI_FORBIDDEN;

                retVal = CryptProtectData(ref plainTextBlob, "", ref entropyBlob, IntPtr.Zero, ref prompt, dwFlags, ref cipherTextBlob);
                if (!retVal)
                    throw new Win32Exception(Marshal.GetLastWin32Error());

                byte[] cipherText = new byte[cipherTextBlob.cbData];
                Marshal.Copy(cipherTextBlob.pbData, cipherText, 0, cipherTextBlob.cbData);
                return cipherText;
            }
            catch (Exception ex)
            {
                throw new DPAPIException("An exception occurred encrypting data", ex);
            }
            finally
            {
                //Always free all allocated memory
                if (plainTextBlob.pbData != IntPtr.Zero)
                    Marshal.FreeHGlobal(plainTextBlob.pbData);

                if (entropyBlob.pbData != IntPtr.Zero)
                    Marshal.FreeHGlobal(entropyBlob.pbData);

                Marshal.FreeHGlobal(cipherTextBlob.pbData);
            }
        }

        public static byte[] Decrypt(byte[] cipherText)
        {
            return Decrypt(cipherText, null);
        }

        public static byte[] Decrypt(byte[] cipherText, byte[] optionalEntropy)
        {
            return Decrypt(cipherText, optionalEntropy, Store.User);
        }

        /// <summary>
        /// Decrypts data using <c>CryptUnprotectData</c>. Returns the clear text data.
        /// </summary>
        /// <param name="cipherText"></param>
        /// <param name="optionalEntropy"></param>
        /// <param name="store"></param>
        /// <returns></returns>
        public static byte[] Decrypt(byte[] cipherText, byte[] optionalEntropy, Store store)
        {
            bool retVal = false;
            DATA_BLOB plainTextBlob = new DATA_BLOB();
            DATA_BLOB cipherTextBlob = new DATA_BLOB();
            DATA_BLOB entropyBlob = new DATA_BLOB();
            CRYPTPROTECT_PROMPTSTRUCT prompt = new CRYPTPROTECT_PROMPTSTRUCT();
            InitPromptstruct(ref prompt);
            int dwFlags;
            try
            {
                int bytesSize = cipherText.Length;
                cipherTextBlob.pbData = Marshal.AllocHGlobal(bytesSize);
                if (cipherTextBlob.pbData == IntPtr.Zero)
                    throw new DPAPIException("Unable to allocate ciphertext buffer");

                cipherTextBlob.cbData = bytesSize;
                Marshal.Copy(cipherText, 0, cipherTextBlob.pbData, bytesSize);

                if (store == Store.Machine)
                {
                    dwFlags = CRYPTPROTECT_LOCAL_MACHINE | CRYPTPROTECT_UI_FORBIDDEN;

                    if (optionalEntropy == null)
                        optionalEntropy = new byte[0];

                    bytesSize = optionalEntropy.Length;
                    entropyBlob.pbData = Marshal.AllocHGlobal(optionalEntropy.Length);
                    if (entropyBlob.pbData == IntPtr.Zero)
                        throw new DPAPIException("Unable to allocate entropy data buffer");

                    Marshal.Copy(optionalEntropy, 0, entropyBlob.pbData, bytesSize);
                    entropyBlob.cbData = bytesSize;
                }
                else
                    dwFlags = CRYPTPROTECT_UI_FORBIDDEN;

                string desc = "";
                retVal = CryptUnprotectData(ref cipherTextBlob, ref desc, ref entropyBlob, IntPtr.Zero, ref prompt, dwFlags, ref plainTextBlob);
                if (!retVal)
                    throw new Win32Exception(Marshal.GetLastWin32Error());

                byte[] plainText = new byte[plainTextBlob.cbData];
                Marshal.Copy(plainTextBlob.pbData, plainText, 0, plainTextBlob.cbData);
                return plainText;
            }
            catch (Exception ex)
            {
                throw new DPAPIException("An exception occurred while decrypting data", ex);
            }
            finally
            {
                //Always free all allocated memory
                if (cipherTextBlob.pbData != IntPtr.Zero)
                    Marshal.FreeHGlobal(cipherTextBlob.pbData);

                if (entropyBlob.pbData != IntPtr.Zero)
                    Marshal.FreeHGlobal(entropyBlob.pbData);

                Marshal.FreeHGlobal(plainTextBlob.pbData);
            }
        }
    }
}
